Add RAFT wrappers around current_device_resource functions #2424
Conversation
harrism
added
feature request
harrism
changed the title
| @@ -375,7 +375,7 @@ rmm::mr::cuda_memory_resource cuda_mr; | |||
| // set the initial size to half of the free device memory | |||
| auto init_size = rmm::percent_of_free_device_memory(50); | |||
| rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{&cuda_mr, init_size}; | |||
| rmm::mr::set_current_device_resource(&pool_mr); // Updates the current device resource pointer to `pool_mr` | |||
| raft::resource::set_current_device_resource(&pool_mr); // Updates the current device resource pointer to `pool_mr` | |||
There was a problem hiding this comment.
Add raft/core/resource/device_memory_resource.hpp to includes of this code example
There was a problem hiding this comment.
I'd suggest to put the new resource helpers somewhere into a detail namespace to discourage the use of these outside of where it is strictly necessary within raft.
Otherwise it goes against raft principle of keeping the program state locally in a raft handle by introducing the global state functions in the api.
In the docs, where we refer to rmm::mr::get_current_device_resource(), I'd rather keep the rmm namespaces, possibly adjusting them to reflect the upcoming API changes in the rmm - again because raft is not supposed to provide this API.
This would severely limit the utility of this PR. I have a dependent PR in each RAPIDS repo that depends on RAFT (cuVS, cuML, cumlprims_mg, cuGraph) to use these RAFT wrappers rather than directly using RMM's functions. |
|
I'm sorry that this limits the utility of the PR, but, in my opinion, this plainly goes against the whole design of raft resources. |
|
Perhaps you can point me to the design document? |
We do outline what @achirkin is talking about in our developer guide. It may sound like something that was done in haste but it was actually very intentional and done for a specific purpose. We've always relied on RMM's APIs to be used outside of RAFT (and outside of cuML, cuGraph, etc...) so that we could merely honor whatever the user has configured with RMM. More recently, we did add a capability for a user to set a different RMM memory resource for temporary memory allocations (and an additionl one for smaller vs larger allocations), but those features still continue to follow the constraint that RAFT maintains no global state of its own. It's dangerous for us to do that because the underlying computational functions now have potential interdependencies that we have to manage and make guarantees with respect to RAFT itself, which now means RAFT needs to be concerned with thread safety. This is a separation of concerns- RAFT, cuML, cuGraph, cuOpt, and others have intentionally avoided doing memory management things in libraries outside of RMM because RMM Is supposed to be responsible for that. If you put this wrapper inside RMM itself then we'd still get the benefits of defining it somewhere upstream (in this case within the RMM namespace) so everyone downstream can use it. Is there any reason we can't just alias this inside of RMM in the meantime? |
|
Aliasing within RMM would just be duplicating something that is already in RMM. The point was to localize changes in non-RMM libraries to a single function rather than have to change every function as we go through this churn. It is painful to modify every function in every RMM-dependent library when RMM changes. |
Fixes #2423
Adds local RAFT wrappers around
rmm::mr::get/set_current_device_resourceandresource_refaliases. The idea is to localize changes needed to keep up with RMM refactoring.The wrappers are added in cpp/include/raft/core/resource/device_memory_resource.hpp
I've marked this PR as breaking because it breaks the ABI, however the API is compatible.